#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
protocol.py - 通信协议实现

Author: xiaobo (xiaobocwb@gmail.com)
Version: 1.0.0
Date: 2024-01

Copyright (c) 2024 xiaobo
Licensed under the MIT License
"""

import serial
import struct
import time
import logging
import threading
from enum import Enum
# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - [Py] %(levelname)s - %(message)s')

# 帧格式常量定义
FRAME_HEAD_MAGIC = 0xAA55   # 帧头
FRAME_HEAD_SIZE = 2         # 帧头长度
FRAME_LENGTH_SIZE = 2       # 长度字段长度
FRAME_MSGID_SIZE = 1        # 消息ID字段长度
FRAME_SEQ_SIZE = 1          # 序列号字段长度
FRAME_CRC_SIZE = 2          # CRC长度

MAX_PAYLOAD_LEN = 1024      # 负载最大长度    
MAX_SEQ = 255               # 序列号最大值

FRAME_BASE_SIZE = FRAME_HEAD_SIZE + FRAME_LENGTH_SIZE + FRAME_MSGID_SIZE + FRAME_SEQ_SIZE + FRAME_CRC_SIZE
MAX_FRAME_SIZE = FRAME_BASE_SIZE + MAX_PAYLOAD_LEN

FRAME_MSGID_ACK = 0xFF      # ACK响应ID


# 解析状态
class ParseState(Enum):
    HEAD1 = 0
    HEAD2 = 1
    LENGTH_L = 2
    LENGTH_H = 3
    MSGID = 4
    SEQ = 5
    DATA = 6
    CRC_L = 7
    CRC_H = 8

class ParseContext:
    def __init__(self):
        self.state = ParseState.HEAD1
        self.data = bytearray()
        self.length = 0
        self.msgid = 0
        self.seq = 0
        self.crc = 0

CRC16_POLY = 0x8005
CRC16_INIT = 0xFFFF



class Protocol:
    def __init__(self, writer):
        """初始化串口协议"""

        self.writer = writer
        self.subscribers = {}  # msgid -> callback
        self.ack_events = {}  # seq -> Event
        self._seq = 0  # 当前序列号
        self.seq_mutex = threading.Lock()
        self.sub_mutex = threading.Lock()
        self.ack_mutex = threading.Lock()

        self.parse_ctx = ParseContext()
        
    def _calc_crc16(self, data):
        crc = CRC16_INIT
        for byte in data:
            crc ^= ((byte << 8) & 0xFFFF)
            for _ in range(8):
                if crc & 0x8000:  # 检查最高位
                    crc = ((crc << 1) & 0xFFFF) ^ CRC16_POLY
                else:
                    crc = (crc << 1) & 0xFFFF
        return crc
    
    def _parse_byte(self, byte):
        parse = self.parse_ctx

        if parse.state == ParseState.HEAD1:
            if byte == (FRAME_HEAD_MAGIC >> 8) & 0xFF:
                parse.state = ParseState.HEAD2

        elif parse.state == ParseState.HEAD2:
            if byte == (FRAME_HEAD_MAGIC & 0xFF):
                parse.state = ParseState.LENGTH_L
                parse.data.clear()
            else:
                parse.state = ParseState.HEAD1

        elif parse.state == ParseState.LENGTH_L:
            parse.length = byte
            parse.data.append(byte)
            parse.state = ParseState.LENGTH_H

        elif parse.state == ParseState.LENGTH_H:
            parse.length |= byte << 8
            parse.data.append(byte)
            min_length = FRAME_SEQ_SIZE + FRAME_MSGID_SIZE + FRAME_CRC_SIZE
            if parse.length < min_length or parse.length > min_length + MAX_PAYLOAD_LEN:
                parse.state = ParseState.HEAD1
            else:
                parse.state = ParseState.MSGID

        elif parse.state == ParseState.MSGID:
            parse.msgid = byte
            parse.data.append(byte)
            parse.length -= 1  # 减去消息ID字节
            parse.state = ParseState.SEQ

        elif parse.state == ParseState.SEQ:
            parse.seq = byte
            parse.data.append(byte)
            parse.length -= 1  # 减去序列号字节
            parse.length -= FRAME_CRC_SIZE  # 减去CRC字节
            if parse.length > 0:
                parse.state = ParseState.DATA
            else:
                parse.state = ParseState.CRC_L

        elif parse.state == ParseState.DATA:
            parse.data.append(byte)
            parse.length -= 1
            if parse.length == 0:
                parse.state = ParseState.CRC_L

        elif parse.state == ParseState.CRC_L:
            parse.crc = byte
            parse.state = ParseState.CRC_H

        elif parse.state == ParseState.CRC_H:
            parse.crc |= byte << 8
            calc_crc = self._calc_crc16(parse.data)
            if parse.crc == calc_crc:
                self._handle_frame()  # 处理完整帧
            else:
                # 校验失败
                logging.warning("[Protocol] CRC check failed, calc: %04X, crc: %04X", calc_crc, parse.crc)
            parse.state = ParseState.HEAD1  # 重置状态

        else:
            parse.state = ParseState.HEAD1

    def parse(self, bytes):
        for byte in bytes:
            self._parse_byte(byte)

    def _handle_frame(self):
        logging.debug(f'_handle_frame msgid:{self.parse_ctx.msgid} data:{self.parse_ctx.data.hex()}')
        if self.parse_ctx.msgid == FRAME_MSGID_ACK:
            # 处理ACK
            with self.ack_mutex:
                seq = self.parse_ctx.data[FRAME_LENGTH_SIZE + FRAME_MSGID_SIZE + FRAME_SEQ_SIZE]
                event = self.ack_events.get(seq & 0x7F)
                if event:
                    logging.debug(f'set event {seq & 0x7F}')
                    event.set()

        else:
            with self.sub_mutex:
                callback = self.subscribers.get(self.parse_ctx.msgid)
                if callback:
                    payload = self.parse_ctx.data[FRAME_LENGTH_SIZE + FRAME_MSGID_SIZE + FRAME_SEQ_SIZE:]
                    callback(self.parse_ctx.msgid, payload)
                    seq = self.parse_ctx.data[FRAME_LENGTH_SIZE + FRAME_MSGID_SIZE]
                    if 0 != (seq & 0x80):
                        # 发送ACK
                        seq &= 0x7F
                        self._send_frame(FRAME_MSGID_ACK, self._alloc_seq(), seq.to_bytes(1, 'little'))

    def _send_frame(self, msgid, seq, payload):
        frame = bytearray()

        # 构建帧
        frame.append((FRAME_HEAD_MAGIC >> 8) & 0xFF)
        frame.append(FRAME_HEAD_MAGIC & 0xFF)

        length = len(payload) + FRAME_MSGID_SIZE + FRAME_SEQ_SIZE + FRAME_CRC_SIZE
        frame.append(length & 0xFF)
        frame.append(length >> 8)

        frame.append(msgid)

        frame.append(seq)

        if payload:
            frame.extend(payload)

        # 计算并添加CRC
        crc = self._calc_crc16(frame[FRAME_HEAD_SIZE:])
        frame.append(crc & 0xFF)
        frame.append((crc >> 8) & 0xFF)

        # 发送帧
        return self.writer(bytes(frame))

    # 分配序列号
    def _alloc_seq(self):
        with self.seq_mutex:
            self._seq = (self._seq + 1) % MAX_SEQ
        return self._seq
       
    def subscribe(self, msgid, callback)->bool:
        """订阅主题
        Args:
            msgid: 消息ID字节(0-254)
            callback: 回调函数，格式为 callback(msgid, payload)
        Returns:
            true 或 false
        """
        if not isinstance(msgid, int) or msgid < 0 or msgid > 254:
            return False
            
        if not callback:
            return False
            
        self.subscribers[msgid] = callback
        logging.info("[Subscribe] Msgid '%d' subscribed locally", msgid)
        return True
        
    def unsubscribe(self, msgid)->bool:
        """取消订阅主题
        Args:
            msgid: 消息ID字节(0-254)
        Returns:
            true 或 false
        """
        if not isinstance(msgid, int) or msgid < 0 or msgid > 254:
            return False
            
        if msgid in self.subscribers:
            del self.subscribers[msgid]
            logging.info("[Unsubscribe] Msgid '%d' unsubscribed locally", msgid)
            
        return True
        
    def publish(self, msgid, payload, need_ack=False, timeout=0.2)->bool:
        """发布消息"""
        if not isinstance(msgid, int) or msgid < 0 or msgid > 254:
            return False
            
        if not isinstance(payload, (bytes, bytearray)) or len(payload) > MAX_PAYLOAD_LEN:
            return False

        for _ in range(255):
            with self.ack_mutex:
                seq = self._alloc_seq()
                is_busy = self.ack_events.get(seq) is not None
                if not is_busy:
                    break

        if is_busy:
            logging.warning("[Protocol] Not found free seq")
            return False

        ack_event = None
        if need_ack:
            seq |= 0x80
            ack_event = threading.Event()
            with self.ack_mutex:
                self.ack_events[seq & 0x7F] = ack_event

        if not self._send_frame(msgid, seq, payload):
            if need_ack:
                with self.ack_mutex:
                    self.ack_events[seq & 0x7F] = None
            return False

        is_ack = True
        if need_ack:
            is_ack = ack_event.wait(timeout=timeout)
            with self.ack_mutex:
                self.ack_events[seq & 0x7F] = None

        return is_ack

    def publish_try(self, msgid, payload, need_ack=False, timeout=0.2, try_times=1)->bool:
        """发布消息，并尝试多次"""
        try_times = try_times if try_times > 0 else 1
        for _ in range(try_times):
            if self.publish(msgid, payload, need_ack, timeout):
                return True
            time.sleep(timeout / 1000)
        return False
